#include <arch/cpu.h>
#include <arch/io.h>
#include <arch/interrupts.h>
#include <arch/vectors.h>
#include <timer.h>
#include <hal/apic.h>
#include <hal/ioapic.h>
#include <sched.h>
#include <process.h>
#include <mm/slab.h>

#define LOG_TAG "TIMER"
#include <log.h>

#define PIT_HZ 1000
#define LATCH (1193180 / PIT_HZ)
#define APIC_TIMER_HZ 1000
#define APIC_TIMER_IC 1000000

static volatile u32 pit_timer_counter = 0;
static volatile u32 apic_timer_done = 1;

volatile u32 sched_counter = 0;

kmem_cache_t* timer_kmem_cache = NULL;

struct timer {
    lnode_t list;
    u32 counter;
    u32 ticks_ms;
    timer_type_t type;
    void* data;
    void (*callback)(void*);
};

lnode_t timer_list = { .next = &timer_list, .prev = &timer_list };

void temp_pit_counter(const isr_param* p) {
    pit_timer_counter++;
}

void temp_apic_handler(const isr_param* p) {
    apic_timer_done--;
    outb(0x21, 0xFF);
}

void nop_isr() {
    // Do nothing
    return;
}

void timer_update() {
    struct timer* timer = NULL, * n;
    list_for_each(timer, n, &timer_list, list) {
        if (timer->counter-- == 0) {
            timer->callback(timer->data);
            if (timer->type == TIMER_ONESHOT) {
                list_remove(&timer->list);
                release_obj(timer_kmem_cache, timer);
            } else {
                timer->counter = timer->ticks_ms;
            }
        }
    }
    sched_counter++;
    if (sched_counter >= TIME_SLICE) {
        sched_counter = 0;
        schedule();
    }
}

static void enable_apic_timer() {
    cli();
    // enable IRQ0 timer interrupt
    outb(0x43, 0x36);		/* binary, mode 3, LSB/MSB, ch 0 */
    outb(0x40, LATCH & 0xFF);	/* LSB */
    outb(0x40, LATCH >> 8);	/* MSB */

    outb(0x21, 0xFE);

    intr_register_isr(PIT_TIMER, temp_pit_counter);
    intr_register_isr(APIC_LINT0_IV, nop_isr);
    intr_register_isr(APIC_TIMER_IV, temp_apic_handler);

    apic_write(APIC_TIMER_LVT, LVT_DELIVERY_FIXED | APIC_TIMER_IV | LVT_TIMER_ONESHOT);
    apic_write(APIC_TIMER_DCR, APIC_TIMER_DIV64);
    apic_write(APIC_TIMER_ICR, APIC_TIMER_IC);

    sti();

    while (apic_timer_done);

    intr_delete_isr(PIT_TIMER, temp_pit_counter);
    intr_delete_isr(APIC_LINT0_IV, nop_isr);

    apic_write(APIC_TIMER_LVT, LVT_DELIVERY_FIXED | APIC_TIMER_IV | LVT_TIMER_PERIODIC);

    intr_register_isr(APIC_TIMER_IV, timer_update);
    if (pit_timer_counter == 0) {
        apic_write(APIC_TIMER_ICR, APIC_TIMER_IC / APIC_TIMER_HZ);
    } else {
        apic_write(APIC_TIMER_ICR, (PIT_HZ / APIC_TIMER_HZ * APIC_TIMER_IC / pit_timer_counter));
    }
    LOGI("APIC Timer initialized\n");
}

static void enable_pit_timer() {
    cli();
    outb(0x43, 0x36);		/* binary, mode 3, LSB/MSB, ch 0 */
    outb(0x40, LATCH & 0xFF);	/* LSB */
    outb(0x40, LATCH >> 8);	/* MSB */
    u8 irq = ioapic_get_irq(0);
    intr_register_isr(PIT_TIMER, timer_update);
    ioapic_redirect(irq, PIT_TIMER, 0, 0);
    sti();
    LOGI("PIT Timer initialized\n");
}

int timer_create(unsigned ticks_ms, timer_type_t type, void (*callback)(void*), void* arg) {
    struct timer* timer = take_obj(timer_kmem_cache);
    if (!timer) {
        return ENOMEM;
    }

    timer->ticks_ms = timer->counter = ticks_ms;
    timer->type = type;
    timer->callback = callback;
    timer->data = arg;
    list_append(&timer_list, &timer->list);
    return 0;
}

void timer_init() {
    if (1) {
        enable_pit_timer();
    } else {
        enable_apic_timer();
    }
    timer_kmem_cache = new_kmem_cache("Timer", sizeof(struct timer), 1, 0);
}